package org.example.weather.bigscreen.core.component;

import com.baomidou.dynamic.datasource.creator.BasicDataSourceCreator;
import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.DataSourceProperty;
import org.apache.ibatis.mapping.Environment;
import org.apache.ibatis.session.SqlSessionFactory;
import org.apache.ibatis.session.SqlSessionFactoryBuilder;
import org.apache.ibatis.transaction.TransactionFactory;
import org.apache.ibatis.transaction.jdbc.JdbcTransactionFactory;
import org.example.weather.bigscreen.core.repository.StWeatherUploadMapper;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;

@Configuration
@EnableConfigurationProperties(ShardingDataSourceProperties.class)
public class ShardingDataSourceConfiguration {
    private final ShardingDataSourceProperties shardingDataSourceProperties;

    public ShardingDataSourceConfiguration(final ShardingDataSourceProperties shardingDataSourceProperties) {
        this.shardingDataSourceProperties = shardingDataSourceProperties;
    }

    @Bean("basicDataSourceMap")
    public Map<String, DataSource> basicDataSourceMap() {
        BasicDataSourceCreator basicDataSourceCreator = new BasicDataSourceCreator();
        Map<String, DataSourceProperty> dataSourcePropertiesMap = shardingDataSourceProperties.getDatasource();
        Map<String, DataSource> dataSourceMap = new HashMap<>(dataSourcePropertiesMap.size() * 2);
        for (Map.Entry<String, DataSourceProperty> item : dataSourcePropertiesMap.entrySet()) {
            String dsName = item.getKey();
            if (!shardingDataSourceProperties.getNames().contains(dsName)) {
                continue;
            }
            DataSourceProperty dataSourceProperty = item.getValue();
            String poolName = dataSourceProperty.getPoolName();
            if (poolName == null || "".equals(poolName)) {
                poolName = dsName;
            }
            dataSourceProperty.setPoolName(poolName);
            dataSourceMap.put(dsName, basicDataSourceCreator.doCreateDataSource(dataSourceProperty));
        }
        return dataSourceMap;
    }

    @Bean("sqlSessionFactoryMap")
    public Map<String, SqlSessionFactory> sqlSessionFactoryMap(@Qualifier("basicDataSourceMap") Map<String, DataSource> basicDataSourceMap) {
        Map<String, SqlSessionFactory> sqlSessionFactoryMap = new HashMap<>(basicDataSourceMap.size());
        basicDataSourceMap.forEach((dataSourceName, dataSource) -> {
            TransactionFactory transactionFactory = new JdbcTransactionFactory();
            Environment environment = new Environment(UUID.randomUUID().toString(), transactionFactory, dataSource);
            org.apache.ibatis.session.Configuration configuration = new org.apache.ibatis.session.Configuration(environment);
            configuration.addMapper(StWeatherUploadMapper.class);
            SqlSessionFactory sqlSessionFactory = new SqlSessionFactoryBuilder().build(configuration);
            sqlSessionFactoryMap.put(dataSourceName, sqlSessionFactory);
        });
        return sqlSessionFactoryMap;
    }
}
